import os

import numpy as np

from .io import write_text_file


def traverse_file(root, lists, filtering, params):
    """
    遍历文件夹

    Parameters
    ----------
    root： 根目录
    lists： 返回的内容
    filtering： 文件过滤器, 定义文件读取规则
    params： 其他参数

    Returns
    -------
    可自定义的信息

    """
    path = os.path.expanduser(root)
    list_files = os.listdir(root)
    for f in list_files:
        file_path = os.path.join(path, f)
        if os.path.isdir(file_path):
            lists = traverse_file(file_path, lists, filtering, params)
        else:
            lists = filtering(file_path, lists, params)
    return lists


def create_train_val_list(file_list, output_root,
                          prefix='', postfix='', train_data_ratio=0.95):
    """
    划分训练集和测试集

    Parameters
    ----------
    file_list: 全体数据的文件列表
    output_root: 保存训练集和测试文件列表的目录
    prefix: 保存文件的前缀
    postfix: 保存文件的后缀
    train_data_ratio: 训练集所占比例

    Returns
    -------

    """
    np.random.shuffle(file_list)
    train = file_list[:int(train_data_ratio * len(file_list))]
    val = file_list[int(train_data_ratio * len(file_list)):]
    train_path = os.path.join(output_root, '{}train{}.txt'.format(prefix, postfix))
    write_text_file(train_path, train)
    val_path = os.path.join(output_root, '{}val{}.txt'.format(prefix, postfix))
    write_text_file(val_path, val)


def create_file_list(data_root,
                     output_file='',
                     filtering=lambda x, y, _: (y[0] + [x], y[1]),
                     params=()):
    """
    遍历一个文件夹， 返回由 `filtering` 定义的文件列表和类别字典

    Parameters
    ----------
    data_root: 文件目录
    output_file: 将文件列表保存的文件名，如果不赋值，则不保存
    filtering: 文件过滤器, 定义文件读取规则
    params: 其他参数

    Returns
    -------

    """
    file_lists = []
    class_num_dict = {}
    file_lists, class_num_dict = traverse_file(
        data_root, [file_lists, class_num_dict], filtering, params)
    if output_file != '':
        write_text_file(output_file, file_lists)
    return file_lists, class_num_dict


def count_instance_number(root, output_file, filtering=lambda x, y, _: (y[0] + [x], y[1]),
                          eliminate=('Others',), split='&!&'):
    """
    统计一个目录下的全部内容的数量，并保存

    Parameters
    ----------
    root: 根目录
    output_file: 保存的文件名
    filtering: 文件过滤器, 定义文件读取规则, 和类别筛选规则
    eliminate: 需要排除的类别
    split: 如果多个内容作为一个条目，可用此分隔符

    Returns
    -------

    """
    _, class_num_dict = create_file_list(root, filtering=filtering, params=([], split))
    l_name = []
    l_num = []
    for key, value in class_num_dict.items():
        if key in eliminate:
            continue
        l_name.append(key)
        l_num.append(value)
    l_name = np.array(l_name)
    l_num = np.array(l_num)
    idx = np.argsort(-l_num)
    l_num = list(map(lambda x: str(x), l_num[idx]))
    l_name = l_name[idx]
    write_text_file(output_file, list(l_num), list(l_name), split=split)
